{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo for `wrf-python`\n",
    "\n",
    "# 1.0 Introduction\n",
    "\n",
    "1) **Created to provide Python with the same functionality found in WRF-NCL**\n",
    "    * WRF variable computation and extraction\n",
    "    * 2D and 3D interpolation routines\n",
    "    * WRF-ARW only\n",
    "    * Metadata\n",
    "2) **Key differences with WRF-NCL**:\n",
    "    * Plotting is handled through matplotlib (with cartopy or basemap), or PyNGL.\n",
    "    * Metadata is optional and can be disabled\n",
    "    \n",
    "**`wrf-python` is a work in progress and should be considered \"pre-alpha\".  Changes to the API are likely to occur between now and the first official release.**    \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# 2.0 Data Structures : `numpy.ndarray` and `xarray.DataArray`\n",
    "\n",
    "\n",
    "`numpy.ndarray`:\n",
    "----------------\n",
    "* Standard array object in Python for array-based computations.\n",
    "* Supports all standard compiled types and also Python types.\n",
    "* Has methods to reshape, perform mathematical calculations, elementwise operations, etc.\n",
    "* Supports missing/fill values using the `numpy.ma.MaskedArray` subclass\n",
    "* Does not support metadata.\n",
    "\n",
    "`xarray.DataArray` (if enabled):\n",
    "--------------------------------\n",
    "* Adds metadata capabilities to the numpy array by wrapping around it.\n",
    "* **HAS A** numpy array, is not a numpy array subclass.\n",
    "* Supports many of the numpy methods, but generally requires the numpy array to be extracted out for computations.\n",
    "* __Always uses NaN to represent missing/fill values__, rather than the MaskedArray type familiar to numpy users.\n",
    "\n",
    "\n",
    "`to_np` method:\n",
    "------------------\n",
    "1. If no missing/fill values, simply extracts the numpy array using the `DataArray.values` property.\n",
    "2. If missing/fill values are found, converts the NaN values to the fill values and returns a MaskedArray."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3.0 Examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1:  Basic Variable Extraction with `getvar` and Print the DataArray and Numpy Array Result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Jupyter Notebook specific command to allow matplotlib images to be shown inline\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from Nio import open_file\n",
    "from wrf import getvar, to_np\n",
    "\n",
    "# Open the output NetCDF file with PyNIO\n",
    "filename = \"wrfout_d01_2010-06-13_21-00-00\"\n",
    "pynio_filename = filename + \".nc\"\n",
    "ncfile = open_file(pynio_filename)\n",
    "\n",
    "# Alternative using netCDF4-python (for reference)\n",
    "# Do 'conda install netcdf4' if the netcdf4 package is not installed\n",
    "#from netCDF4 import Dataset as nc\n",
    "#filename = \"wrfout_d01_2010-06-13_21-00-00\"\n",
    "#ncfile = nc(filename)\n",
    "\n",
    "# Extract the terrain height, which will be returned as an xarray.DataArray object (if available).  DataArray\n",
    "# objects include meta data, similar to NCL's variables.\n",
    "# Note:  can also use the 'ter' variable\n",
    "terrainx = getvar(ncfile, \"HGT\", timeidx=0)\n",
    "\n",
    "print (terrainx)\n",
    "\n",
    "# To extract the numpy array, use the to_np function\n",
    "terrain_numpy = to_np(terrainx)\n",
    "print (\"\\nExtracted numpy array:\\n\")\n",
    "print (terrain_numpy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2:  Use `getvar` to Extract the 'HGT' Variable and Make a Plot Using Matplotlib with Basemap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "from Nio import open_file\n",
    "from wrf import getvar, to_np, get_basemap, latlon_coords\n",
    "\n",
    "# Open the output NetCDF file with PyNIO\n",
    "filename = \"wrfout_d01_2010-06-13_21-00-00\"\n",
    "pynio_filename = filename + \".nc\"\n",
    "ncfile = open_file(pynio_filename)\n",
    "\n",
    "# Extract the terrain height, which will be returned as an xarray.DataArray object (if available).  DataArray\n",
    "# objects include meta data, similar to NCL's variables.\n",
    "# Note:  can also use the 'ter' variable\n",
    "terrainx = getvar(ncfile, \"HGT\", timeidx=0)\n",
    "\n",
    "# Use to_np to extract the numpy array, since matplotlib does not handle xarray.DataArray natively\n",
    "terrain_data = to_np(terrainx)\n",
    "\n",
    "# Get the lat/lon 2D coordinate arrays.  Use to_np to extract the numpy array since basemap does not\n",
    "# handle xarray.DataArray natively.\n",
    "lons = latlon_coords(terrainx, as_np=True)\n",
    "lats = latlon_coords(terrainx, as_np=True)\n",
    "\n",
    "# Extract the basemap object from the projection information\n",
    "bm = get_basemap(terrainx)\n",
    "\n",
    "# Convert the lat/lon coordinates to projected x,y\n",
    "x,y = bm(lons, lats)\n",
    "\n",
    "# Create the figure\n",
    "fig = plt.figure(figsize=(16,16))\n",
    "ax = fig.add_axes([0.1,0.1,0.8,0.8])\n",
    "\n",
    "# Draw filled contours from 100 to 3000 m, every 200 meters.\n",
    "levels = np.arange(100, 3000, 200)\n",
    "bm.contourf(x, y, terrain_data, levels=levels, extend=\"max\", cmap=get_cmap(\"terrain\"))\n",
    "\n",
    "# Draw the coastlines and country borders.\n",
    "bm.drawcoastlines()\n",
    "bm.drawcountries()\n",
    "\n",
    "# Draw the color bar\n",
    "plt.colorbar(ax=ax, shrink=.7)\n",
    "\n",
    "# Add a title\n",
    "plt.title(\"Terrain Height (m)\", {\"fontsize\" : 20})\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 3:  Use `getvar` to Compute Dewpoint and Make a Plot Using Matplotlib with Basemap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "from Nio import open_file\n",
    "from wrf import getvar, to_np, get_basemap, latlon_coords\n",
    "\n",
    "# Open the output NetCDF file with PyNIO\n",
    "filename = \"wrfout_d01_2010-06-13_21-00-00\"\n",
    "pynio_filename = filename + \".nc\"\n",
    "ncfile = open_file(pynio_filename)\n",
    "\n",
    "# Extract the dewpoint, which will be returned as an xarray.DataArray object (if available).  DataArray\n",
    "# objects include meta data, similar to NCL's variables.\n",
    "dewpointx = getvar(ncfile, \"td\", timeidx=0)\n",
    "\n",
    "\n",
    "# Dewpoint is a 3D variable, so let's just use the lowest level\n",
    "dewpointx_sfc = dewpointx[0,:,:]\n",
    "\n",
    "# Use to_np to extract the numpy array, since matplotlib does not handle xarray.DataArray natively\n",
    "dewpoint_ndarray = to_np(dewpointx_sfc)\n",
    "\n",
    "# Get the lat/lon 2D coordinate arrays.  Use to_np to extract the numpy array since basemap does not\n",
    "# handle xarray.DataArray natively.\n",
    "lons = latlon_coords(dewpointx_sfc, as_np=True)\n",
    "lats = latlon_coords(dewpointx_sfc, as_np=True)\n",
    "\n",
    "# Extract the basemap object from the projection information\n",
    "bm = get_basemap(dewpointx_sfc)\n",
    "\n",
    "# Convert the lat/lon coordinates to projected x,y\n",
    "x,y = bm(lons, lats)\n",
    "\n",
    "# Create the figure\n",
    "fig = plt.figure(figsize=(16,16))\n",
    "ax = fig.add_axes([0.1,0.1,0.8,0.8])\n",
    "\n",
    "# Draw filled contours from -20 C to 40 C, every 5 C.\n",
    "levels = np.arange(-20, 40, 5)\n",
    "bm.contourf(x, y, dewpoint_ndarray, levels=levels, extend=\"both\", cmap=get_cmap(\"RdYlGn\"))\n",
    "\n",
    "# Draw the coastlines, country borders, and states.\n",
    "bm.drawcoastlines()\n",
    "bm.drawcountries()\n",
    "bm.drawstates()\n",
    "\n",
    "# Draw the color bar\n",
    "plt.colorbar(ax=ax, shrink=.7)\n",
    "\n",
    "# Add a title\n",
    "plt.title(\"Dewpoint Temperature (C)\", {\"fontsize\" : 20})\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 4: Create a Vertical Cross-Section Using `vertcross` and Make a Plot Using Matplotlib with Basemap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "from Nio import open_file\n",
    "from wrf import getvar, vertcross, to_np\n",
    "\n",
    "# Open the output NetCDF file with PyNIO\n",
    "filename = \"wrfout_d01_2010-06-13_21-00-00\"\n",
    "pynio_filename = filename + \".nc\"\n",
    "ncfile = open_file(pynio_filename)\n",
    "\n",
    "# Extract pressure and model height\n",
    "p = getvar(ncfile, \"pressure\", timeidx=0)\n",
    "z = getvar(ncfile, \"z\", timeidx=0)\n",
    "\n",
    "# Define a horizontal cross section going left to right using a pivot point in the center of the grid\n",
    "# with an angle of 90.0 degrees (0 degrees is vertical).\n",
    "# Pivot point is a tuple of (x, y)\n",
    "pivot_point = (z.shape[-1] / 2, z.shape[-2] / 2) \n",
    "angle = 90.0\n",
    "\n",
    "# Compute the vertical cross-section interpolation.  Also, include the lat/lon points along the cross-section.\n",
    "p_vertx = vertcross(p, z, pivot_point=pivot_point, angle=angle, latlon=True)\n",
    "\n",
    "# Extract the numpy array\n",
    "p_vert_array = to_np(p_vertx)\n",
    "\n",
    "# Create the figure\n",
    "fig = plt.figure(figsize=(20,8))\n",
    "ax = plt.axes([0.1,0.1,0.8,0.8])\n",
    "\n",
    "# Define the levels [0, 50, 100, 150, ....]\n",
    "levels = [0 + 50*n for n in range(20)]\n",
    "\n",
    "# Make the contour plot\n",
    "plt.contour(p_vert_array, levels=levels)\n",
    "\n",
    "# Add the color bar\n",
    "plt.colorbar(ax=ax)\n",
    "\n",
    "# Set the x-ticks to use latitude and longitude labels.\n",
    "coord_pairs = to_np(p_vertx.coords[\"xy_loc\"])\n",
    "x_ticks = np.arange(coord_pairs.shape[0])\n",
    "x_labels = [pair.latlon_str() for pair in to_np(coord_pairs)]\n",
    "plt.xticks(x_ticks[::100], x_labels[::100]) # Only use every 100th tick.\n",
    "\n",
    "# Set the y-ticks to be height.\n",
    "vert_vals = to_np(p_vertx.coords[\"vertical\"])\n",
    "v_ticks = np.arange(vert_vals.shape[0])\n",
    "plt.yticks(v_ticks[::10], vert_vals[::10]) # Only use every 10th tick.\n",
    "\n",
    "# Set the x-axis and  y-axis labels\n",
    "plt.xlabel(\"Latitude, Longitude\", fontsize=14)\n",
    "plt.ylabel(\"Height (m)\", fontsize=14)\n",
    "\n",
    "# Add a title\n",
    "plt.title(\"Vertical Cross-Section of Pressure (hPa)\", {\"fontsize\" : 20})\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 5: Interpolate Wind and Height to the 500 hPa Level Using `interplevel` and Make a Plot of the 500 hPa Wind Barbs, Wind Speed, and Height Contours Using Matplotlib with Basemap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "from Nio import open_file\n",
    "from wrf import getvar, interplevel, to_np, get_basemap, latlon_coords\n",
    "\n",
    "# Open the output NetCDF file with PyNIO\n",
    "filename = \"wrfout_d01_2010-06-13_21-00-00\"\n",
    "pynio_filename = filename + \".nc\"\n",
    "ncfile = open_file(pynio_filename)\n",
    "\n",
    "# Extract pressure, model height, u and v winds on mass points\n",
    "p = getvar(ncfile, \"pressure\", timeidx=0)\n",
    "z = getvar(ncfile, \"z\", timeidx=0, units=\"dm\")\n",
    "ua = getvar(ncfile, \"ua\", timeidx=0, units=\"kts\")\n",
    "va = getvar(ncfile, \"va\", timeidx=0, units=\"kts\")\n",
    "wspd = getvar(ncfile, \"wspd_wdir\", timeidx=0, units=\"kts\")[0,...]\n",
    "\n",
    "# Interpolate height, u, and v to to 500 hPa\n",
    "ht_500 = interplevel(z, p, 500)\n",
    "u_500 = interplevel(ua, p, 500)\n",
    "v_500 = interplevel(va, p, 500)\n",
    "wspd_500 = interplevel(wspd, p, 500)\n",
    "\n",
    "# Get the projection\n",
    "bm = get_basemap(p)\n",
    "\n",
    "# Basemap needs numpy arrays, extract with to_np\n",
    "lons = latlon_coords(ht_500, as_np=True)\n",
    "lats = latlon_coords(ht_500, as_np=True)\n",
    "\n",
    "# Convert the lat/lon coordinates to projected x,y\n",
    "x,y = bm(lons, lats)\n",
    "\n",
    "fig = plt.figure(figsize=(20,20))\n",
    "ax = plt.axes([0.1,0.1,0.8,0.8])\n",
    "\n",
    "# Draw the coastlines and country borders.\n",
    "bm.drawcoastlines()\n",
    "bm.drawcountries()\n",
    "bm.drawstates()\n",
    "\n",
    "# Make the 500 hPa height contours\n",
    "ht_contours = bm.contour(x, y, to_np(ht_500), 10, linewidths=2.0, colors=\"black\")\n",
    "\n",
    "# Use contour labels for height\n",
    "plt.clabel(ht_contours, inline=True, fontsize=12, fmt=\"%i\")\n",
    "\n",
    "# Make the wind speed filled contours\n",
    "levels = np.arange(40, 120, 10)\n",
    "bm.contourf(x, y, to_np(wspd_500), levels=levels, extend=\"max\", cmap=get_cmap(\"rainbow\"))\n",
    "\n",
    "# Make the wind barbs.  Only use every 50th in each direction.\n",
    "bm.barbs(x[::50,::50], y[::50,::50], to_np(u_500[::50, ::50]), to_np(v_500[::50, ::50]))\n",
    "\n",
    "# Make the color bar\n",
    "plt.colorbar(ax=ax, shrink=.7)\n",
    "\n",
    "# Make the title\n",
    "plt.title(\"500 MB Heights (dm), Wind Speed (kts), and Wind Barbs (kts)\", {\"fontsize\" : 20})\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
